import submitit
import argparse, random

from ecgcmr.signal.get_embeddings import get_embeddings
from ecgcmr.utils.distributed import init_dist_node


def parse_args():
	parser = argparse.ArgumentParser()
	parser.add_argument('-slurm', action='store_true', help='Submit with slurm')
	parser.add_argument('-slurm_ngpus', type=int, default = 2, help='num of gpus per node')
	parser.add_argument('-slurm_nnodes', type=int, default = 1, help='number of nodes')
	parser.add_argument('-slurm_nodelist', default = None, help='slurm nodeslist')
	parser.add_argument('-slurm_partition', type=str, default = "general", help='slurm partition')
	parser.add_argument('-slurm_timeout', type=int, default = 2800, help='slurm timeout minimum, reduce if running on the "Quick" partition')
	parser.add_argument('-output_dir', type=str, help='path to shared output directory')


	parser.add_argument('-ecg_data_path', type=str, default="/vol/space/projects/ukbb/projects/ecg/ecgs_train_ecg_imaging_noBase_gn.pt")
	parser.add_argument('-checkpoint_path', type=str, default="TIME-ECG-CMR/MMCL-ECG-CMR/model_weights/signal_encoder_mdm.pth")
	parser.add_argument('-ecg_time_steps', type=int, default=2500)
	parser.add_argument('-input_electrodes', type=int, default=12)
	parser.add_argument('-patch_width', type=int, default=1)
	parser.add_argument('-patch_height', type=int, default=100)

	parser.add_argument('-batch_size', type=int, default=8)
	parser.add_argument('-num_workers', type=int, default=2)

	args = parser.parse_args()

	return args


def main(args):
	
	args.port = random.randint(49152,65535)

	class Slurm_Evalute(object):
		def __init__(self, args):
			self.args = args

		def __call__(self):

			init_dist_node(self.args)
			
			get_embeddings(None, self.args)

	executor = submitit.AutoExecutor(folder="submitit_logs")
	executor.update_parameters(
		mem_gb=40*args.slurm_ngpus,
		gpus_per_node=args.slurm_ngpus,
		tasks_per_node=args.slurm_ngpus,
		cpus_per_task=10,
		nodes=args.slurm_nnodes,
		slurm_exclude="prometheus,atlas",
		timeout_min=args.slurm_timeout,
		slurm_partition="normal",
	)

	trainer = Slurm_Evalute(args)
	job = executor.submit(trainer)

	print(f"Job submitted, ID: {job.job_id}")

if __name__ == '__main__':
	args = parse_args()
	main(args)
